import cv2
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image

# -----------------------------
# 修正后的标定参数（假设主点已修正）
# -----------------------------
mtx_main = np.array([[1093.23802, 0, 635.746975],  # 假设修正后的合理主点
                     [0, 1096.467663, 337.6218325],
                     [0, 0, 1]], dtype=np.float64)

dist_main = np.array([-0.149254721, 0.54330647, 0, 0, 0], dtype=np.float64).reshape(1, -1)

mtx_side = np.array([[1091.510188, 0, 659.0351622],
                     [0, 1093.072199, 298.8419753],
                     [0, 0, 1]], dtype=np.float64)
dist_side = np.array([-0.034262054, -0.123556241, 0, 0, 0], dtype=np.float64).reshape(1, -1)

frame_size = (1280, 720)

R = np.array([[0.999296745, -0.016515227, 0.033663969],
              [0.0155981, 0.999504858, 0.027326502],
              [-0.034098604, -0.026782191, 0.999059558]], dtype=np.float64)
T = np.array([-89.42851336, -0.510686004, -0.285366743], dtype=np.float64).reshape(3, 1)

# -----------------------------
# 立体校正（调整alpha=0.5）
# -----------------------------
R1, R2, P1, P2, Q, roi1, roi2 = cv2.stereoRectify(
    mtx_main, dist_main, mtx_side, dist_side,
    frame_size, R, T,
    flags=cv2.CALIB_ZERO_DISPARITY,
    alpha=0.5  # 调整此处
)

# -----------------------------
# 图像校正（根据roi裁剪）
# -----------------------------
def rectify_image(img, mtx, dist, R, P):
    h, w = img.shape[:2]
    new_mtx = P[:3, :3]
    mapx, mapy = cv2.initUndistortRectifyMap(mtx, dist, R, new_mtx, (w, h), cv2.CV_32FC1)
    rectified = cv2.remap(img, mapx, mapy, cv2.INTER_LINEAR)
    return rectified

img_main = cv2.imread('data1/left01.jpg')
img_side = cv2.imread('data1/right01.jpg')

rectified_main = rectify_image(img_main, mtx_main, dist_main, R1, P1)
rectified_side = rectify_image(img_side, mtx_side, dist_side, R2, P2)

# 根据roi裁剪
x1, y1, w1, h1 = roi1
x2, y2, w2, h2 = roi2
rectified_main = rectified_main[y1:y1+h1, x1:x1+w1]
rectified_side = rectified_side[y2:y2+h2, x2:x2+w2]

# -----------------------------
# 绘制极线（更密集的线条）
# -----------------------------
rectified_main_rgb = cv2.cvtColor(rectified_main, cv2.COLOR_BGR2RGB)
rectified_side_rgb = cv2.cvtColor(rectified_side, cv2.COLOR_BGR2RGB)

im_L = Image.fromarray(rectified_main_rgb)
im_R = Image.fromarray(rectified_side_rgb)

width_total = im_L.size[0] + im_R.size[0]
height = im_L.size[1]
img_compare = Image.new('RGB', (width_total, height))
img_compare.paste(im_L, box=(0, 0))
img_compare.paste(im_R, box=(im_L.size[0], 0))

plt.figure(figsize=(20, 20))
plt.imshow(img_compare)

num_lines = 20  # 增加线条密度
line_spacing = height / num_lines
for i in range(1, num_lines):
    y = i * line_spacing
    plt.axhline(y=y, color='r', linestyle='-', linewidth=0.5)

plt.axis('off')
plt.savefig('epipolar_lines.png', bbox_inches='tight', pad_inches=0)
plt.show()